import numpy as np
import matplotlib.pyplot as plt
# 设置matplotlib支持中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体显示中文
plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号

class CommunicationSystem:
    def __init__(self, frequency1=10, frequency2=20, sampling_rate=1000, symbol_duration=0.1, snr_db=10):
        self.frequency1 = frequency1  # 0对应的频率
        self.frequency2 = frequency2  # 1对应的频率
        self.sampling_rate = sampling_rate  # 采样率
        self.symbol_duration = symbol_duration  # 每个符号的持续时间
        self.snr_db = snr_db  # 信噪比(dB)
        self.samples_per_symbol = int(sampling_rate * symbol_duration)  # 每个符号的采样点数
    
    def text_to_binary(self, text):
        """将文本转换为二进制编码（使用UTF-8）"""
        # 先转换为bytes，再转换为二进制字符串
        byte_data = text.encode('utf-8')
        binary = ''.join(format(byte, '08b') for byte in byte_data)
        return binary
    
    def binary_to_text(self, binary):
        """将二进制编码转换回文本（使用UTF-8）"""
        # 确保二进制字符串长度是8的倍数
        while len(binary) % 8 != 0:
            binary = '0' + binary
        
        # 先转换为bytes
        byte_list = []
        for i in range(0, len(binary), 8):
            byte = binary[i:i+8]
            byte_list.append(int(byte, 2))
        
        # 再解码为文本
        try:
            text = bytes(byte_list).decode('utf-8')
            return text
        except UnicodeDecodeError:
            # 如果解码失败，返回原始字节的十六进制表示
            return f"解码失败: {bytes(byte_list).hex()}"
    
    def modulate_2fsk(self, binary):
        """使用2FSK调制二进制信号"""
        time = np.linspace(0, len(binary) * self.symbol_duration, len(binary) * self.samples_per_symbol)
        modulated_signal = np.zeros_like(time)
        
        for i, bit in enumerate(binary):
            start_idx = i * self.samples_per_symbol
            end_idx = (i + 1) * self.samples_per_symbol
            t = time[start_idx:end_idx] - i * self.symbol_duration
            
            if bit == '0':
                modulated_signal[start_idx:end_idx] = np.sin(2 * np.pi * self.frequency1 * t)
            else:
                modulated_signal[start_idx:end_idx] = np.sin(2 * np.pi * self.frequency2 * t)
        
        return time, modulated_signal
    
    def add_noise(self, signal):
        """向信号添加高斯白噪声"""
        snr = 10 ** (self.snr_db / 10)
        signal_power = np.mean(signal ** 2)
        noise_power = signal_power / snr
        noise = np.sqrt(noise_power) * np.random.randn(len(signal))
        return signal + noise
    
    def demodulate_2fsk(self, received_signal, binary_length):
        """解调2FSK信号"""
        recovered_bits = ''
        
        for i in range(binary_length):
            start_idx = i * self.samples_per_symbol
            end_idx = (i + 1) * self.samples_per_symbol
            segment = received_signal[start_idx:end_idx]
            
            # 创建参考信号
            t = np.linspace(0, self.symbol_duration, self.samples_per_symbol)
            ref1 = np.sin(2 * np.pi * self.frequency1 * t)
            ref2 = np.sin(2 * np.pi * self.frequency2 * t)
            
            # 相关解调
            corr1 = np.sum(segment * ref1)
            corr2 = np.sum(segment * ref2)
            
            # 判决
            if corr1 > corr2:
                recovered_bits += '0'
            else:
                recovered_bits += '1'
        
        return recovered_bits
    
    def plot_signals(self, time, original_signal, modulated_signal, noisy_signal, demodulated_binary):
        """绘制信号波形图"""
        fig, axs = plt.subplots(4, 1, figsize=(12, 10))
        
        # 绘制原始二进制信号
        axs[0].step(np.arange(len(original_signal)) * self.symbol_duration, 
                   [int(bit) for bit in original_signal], where='post')
        axs[0].set_title('原始二进制信号')
        axs[0].set_xlabel('时间 (s)')
        axs[0].set_ylabel('幅度')
        axs[0].set_ylim(-0.5, 1.5)
        axs[0].grid(True)
        
        # 绘制调制信号
        axs[1].plot(time[:1000], modulated_signal[:1000])  # 只显示前1000个采样点以便观察
        axs[1].set_title('2FSK调制信号')
        axs[1].set_xlabel('时间 (s)')
        axs[1].set_ylabel('幅度')
        axs[1].grid(True)
        
        # 绘制带噪声的信号
        axs[2].plot(time[:1000], noisy_signal[:1000])  # 只显示前1000个采样点以便观察
        axs[2].set_title('带噪声的2FSK信号')
        axs[2].set_xlabel('时间 (s)')
        axs[2].set_ylabel('幅度')
        axs[2].grid(True)
        
        # 绘制解调后的二进制信号
        axs[3].step(np.arange(len(demodulated_binary)) * self.symbol_duration, 
                   [int(bit) for bit in demodulated_binary], where='post')
        axs[3].set_title('解调后的二进制信号')
        axs[3].set_xlabel('时间 (s)')
        axs[3].set_ylabel('幅度')
        axs[3].set_ylim(-0.5, 1.5)
        axs[3].grid(True)
        
        plt.tight_layout()
        plt.show()
    
    def run(self, input_text):
        """运行完整的通信系统流程"""
        print(f"输入文本: {input_text}")
        
        # 1. 信源转二进制
        binary = self.text_to_binary(input_text)
        print(f"二进制编码: {binary}")
        print(f"二进制长度: {len(binary)} bits")
        
        # 2. 调制
        time, modulated_signal = self.modulate_2fsk(binary)
        print(f"调制后信号长度: {len(modulated_signal)} 采样点")
        
        # 3. 添加噪声（模拟信道）
        noisy_signal = self.add_noise(modulated_signal)
        
        # 4. 解调
        demodulated_binary = self.demodulate_2fsk(noisy_signal, len(binary))
        print(f"解调后二进制: {demodulated_binary}")
        
        # 5. 二进制转文本
        recovered_text = self.binary_to_text(demodulated_binary)
        print(f"解码后文本: {recovered_text}")
        
        # 6. 计算误码率
        bit_errors = sum(1 for a, b in zip(binary, demodulated_binary) if a != b)
        ber = bit_errors / len(binary) if len(binary) > 0 else 0
        print(f"误码率: {ber:.6f}")
        
        # 检查传输是否成功
        transmission_success = input_text == recovered_text
        print(f"通信{'成功' if transmission_success else '失败'}")
        
        # 7. 绘制信号波形
        self.plot_signals(time, binary, modulated_signal, noisy_signal, demodulated_binary)
        
        return recovered_text, ber

# 测试程序
if __name__ == "__main__":
    # 创建通信系统实例
    # 增加信噪比，降低误码率
    comm_system = CommunicationSystem(frequency1=5, frequency2=15, sampling_rate=1000, symbol_duration=0.1, snr_db=20)
    
    # 输入文本
    user_input = input("请输入要传输的文本: ")
    if not user_input:
        user_input = "Hello, 通信系统!"  # 默认输入
    
    # 运行系统
    recovered_text, ber = comm_system.run(user_input)